#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_matmul.h"
#include "common.h"
#include "matrix_mul.h"

using namespace common;

namespace matrix {

int Mul(aclrtStream *stream, float* x, float* y, float* r, int m, int k, int n) {
  // 1、创建Tensor
  std::vector<int64_t> xShape = {m, k};
  std::vector<int64_t> yShape = {k ,n};
  std::vector<int64_t> rShape = {m, n};

  void* xDeviceAddr = nullptr;
  void* yDeviceAddr = nullptr;
  void* rDeviceAddr = nullptr;

  aclTensor* xTensor = nullptr;
  aclTensor* yTensor = nullptr;
  aclTensor* rTensor = nullptr;

  
  auto ret = CreateAclTensor(x, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &xTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  ret = CreateAclTensor(y, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &yTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  ret = CreateAclTensor(r, rShape, &rDeviceAddr, aclDataType::ACL_FLOAT, &rTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 2、调用CANN算子库API
  int8_t cubeMathType = 1;
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor = nullptr;
  // 调用aclnnMatmul第一段接口
  ret = aclnnMatmulGetWorkspaceSize(xTensor, yTensor, rTensor, cubeMathType, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmulGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // 根据第一段接口计算出的workspaceSize申请device内存
  void* workspaceAddr = nullptr;
  if (workspaceSize > 0) {
    ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
  }
  // 调用aclnnMatmul第二段接口
  ret = aclnnMatmul(workspaceAddr, workspaceSize, executor, *stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMatmul failed. ERROR: %d\n", ret); return ret);

  // 3、同步等待任务执行结束
  ret = aclrtSynchronizeStream(*stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 4、将device侧内存上的结果拷贝至host侧
  ret = aclrtMemcpy(r, m * n * sizeof(float), rDeviceAddr, m * n * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);

  // 5. 释放aclTensor和aclScalar
  aclDestroyTensor(xTensor);
  aclDestroyTensor(yTensor);
  aclDestroyTensor(rTensor);

  // 6、释放device资源
  aclrtFree(xDeviceAddr);
  aclrtFree(yDeviceAddr);
  aclrtFree(rDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  return 0;
}

}